# 模型训练

import paddle
from paddleseg.models import DeepLabV3, UNetPlusPlus
from paddleseg.models.backbones import ResNet18_vd, ResNet34_vd, ResNet50_vd, ResNet101_vd, ResNet152_vd
from paddleseg.models.losses import CrossEntropyLoss
from src.config.config import NUM_CLASSES, EXP_DIR, ITERS, BACH_SIZE, SAVE_ITERS, LOG_ITERS, NUM_WORKERS, PRETRAINED_URL
from src.util.dataset import train_dataset, val_dataset

# pretrained参数，使用paddle预训练的deeplabv3_resnet50模型
model = DeepLabV3(
    num_classes=NUM_CLASSES,
    backbone=ResNet50_vd(),  # currently support Resnet50_vd/Resnet101_vd/Xception65.
    pretrained=PRETRAINED_URL
)
# model = UNetPlusPlus(in_channels=3, num_classes=NUM_CLASSES, use_deconv=True)

# 设置学习率、优化器
base_lr = 0.01
# base_lr = 0.001
lr = paddle.optimizer.lr.PolynomialDecay(base_lr, power=0.9, decay_steps=3000, end_lr=0)
optimizer = paddle.optimizer.SGD(lr, parameters=model.parameters(), weight_decay=4.0e-5)

# 设置损失函数
losses = {}
losses['types'] = [CrossEntropyLoss()] * 1
losses['coef'] = [1] * 1

# 开始训练
from paddleseg.core import train

train(
    model=model,
    train_dataset=train_dataset,
    val_dataset=val_dataset,
    optimizer=optimizer,
    losses=losses,
    save_dir=EXP_DIR,  # 模型和visualdl日志文件的保存根路径
    iters=ITERS,  # 训练迭代次数
    batch_size=BACH_SIZE,  # 单卡batch size
    save_interval=SAVE_ITERS,  # 模型保存的间隔步数
    log_iters=LOG_ITERS,  # 打印日志的间隔步数
    num_workers=NUM_WORKERS,  # 用于异步读取数据的进程数量， 大于等于1时开启子进程读取数据
    use_vdl=True,  # 是否开启visualdl记录训练数据
)
